import torch
import torch.nn as nn
from vortex.network.head.yolov3 import ConvBnLeaky, ConvBnLeakyEmbedding
from vortex.network.backbone import darknet53


class YoloBody(nn.Module):
    def __init__(self, anchors, num_classes):
        super(YoloBody, self).__init__()
        """
        0 for deepest output.
        """
        self.anchors = anchors  
        self.num_classes = num_classes

        self.backbone = darknet53(pretrained=False)
        backbone_out0 = 1024
        backbone_out1 = 512
        backbone_out2 = 256

        filters0 = len(self.anchors[0]) * (5 + self.num_classes)
        filters1 = len(self.anchors[1]) * (5 + self.num_classes)
        filters2 = len(self.anchors[2]) * (5 + self.num_classes)
        
        self.embedding0 = ConvBnLeakyEmbedding([512, 1024], backbone_out0, filters0)
        self.cbl1 = ConvBnLeaky(512, 256, 1)
        self.upsample1 = nn.Upsample(scale_factor=2, mode='nearest')
        self.embedding1 = ConvBnLeakyEmbedding([256, 512], backbone_out1 + 256, filters1)
        self.cbl2 = ConvBnLeaky(256, 128, 1)
        self.upsample2 = nn.Upsample(scale_factor=2, mode='nearest')
        self.embedding2 = ConvBnLeakyEmbedding([128, 256], backbone_out2 + 128, filters2)

    def forward(self, x):
        x2, x1, x0 = self.backbone(x)

        out0, out0_branch = self.embedding0(x0)
        out0_branch = self.cbl1(out0_branch)
        out0_branch = self.upsample1(out0_branch)
        x1 = torch.cat([out0_branch, x1], 1)
        out1, out1_branch = self.embedding1(x1)
        out1_branch = self.cbl2(out1_branch)
        out1_branch = self.upsample2(out1_branch)
        x2 = torch.cat([out1_branch, x2], 1)
        out2, _ = self.embedding2(x2)

        return out0, out1, out2
